package com.ccut.xunfeiSpringboot.utils;

import cn.hutool.json.JSONUtil;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;

import java.io.IOException;
import java.net.ProtocolException;
import java.util.Objects;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

@Slf4j
public class AvatarWsUtil {
    public WebSocket webSocket;
    private static final AtomicBoolean isConnected = new AtomicBoolean(false);
    private static CountDownLatch countDownLatch;
    private static CountDownLatch connect;

    // 新增消息回调函数
    private Consumer<String> messageCallback;

    // 设置消息回调
    public void setMessageCallback(Consumer<String> callback) {
        this.messageCallback = callback;
    }

    public AvatarWsUtil(String requestUrl) {
        Request wsRequest = (new Request.Builder()).url(requestUrl).build();
        OkHttpClient okHttpClient = new OkHttpClient().newBuilder().build();
        connect = new CountDownLatch(1);
        this.webSocket = okHttpClient.newWebSocket(wsRequest, buildWebSocketListener());
    }

    public void start(JSONObject request, CountDownLatch countDownLatch) throws Exception {
        this.countDownLatch = countDownLatch;
        connect.await();
        send(request);
    }

    public void send(JSONObject request) {
        if (isConnected.get()) {
            String jsonStr = JSONUtil.toJsonStr(request);
            log.info("send :{}", jsonStr);
            webSocket.send(jsonStr);
        }
    }

    private WebSocketListener buildWebSocketListener() {
        return new WebSocketListener() {
            public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
                log.info("onOpen");
                isConnected.set(true);
                connect.countDown();
            }

            public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
                log.info("onMessage: {}", text);

                // 调用消息回调函数
                if (messageCallback != null) {
                    messageCallback.accept(text);
                }

                JSONObject jsonObject = JSON.parseObject(text);
                int code = jsonObject.getJSONObject("header").getIntValue("code");
                if (code != 0) {
                    onEvent(webSocket, 1002, jsonObject.getJSONObject("header").getString("message"), "server closed");
                    return;
                }
                JSONObject payload = jsonObject.getJSONObject("payload");
                if (payload != null) {
                    JSONObject avatar = payload.getJSONObject("avatar");
                    if (avatar != null && avatar.getString("stream_url") != null &&
                            !avatar.getString("stream_url").isEmpty()) {
                        countDownLatch.countDown();
                    }
                }
            }

            public void onClosing(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
                this.onEvent(webSocket, code, reason, "onClosing");
            }

            public void onClosed(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
                this.onEvent(webSocket, code, reason, "onClosed");
            }

            public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable tx, Response response) {
                Object t;
                try {
                    String responseBody = response != null && response.body() != null ?
                            response.body().string() : "No response body";
                    JSONObject body = JSON.parseObject(responseBody);
                    t = new ProtocolException(body.toString());
                } catch (IOException var6) {
                    t = var6;
                }
                log.info("onFailure:{}", t);
            }

            void onEvent(@NotNull WebSocket webSocket, int code, String reason, String event) {
                log.info("session {} . code:{}, reason:{}", event, code, reason);
                isConnected.set(false);
                if (countDownLatch != null) {
                    countDownLatch.countDown();
                }
                try {
                    webSocket.close(code, reason);
                } catch (Exception var6) {
                    log.error("{} error.{}", event, var6.getMessage());
                }
            }
        };
    }

    public void close() {
        if (webSocket != null) {
            this.webSocket.close(1000, "");
        }
    }
}